#include <gtest/gtest.h>
#include "sandbox.h"
#include <chrono>
#include <thread>
#include <filesystem>
#include <fstream>
#include <spdlog/spdlog.h>

using namespace judge;
using namespace std::chrono_literals;

class SandboxTest : public ::testing::Test {
protected:
    void SetUp() override {
        // 设置日志级别为 error，减少输出
        spdlog::set_level(spdlog::level::err);
        
        // 确保工作目录存在
        std::filesystem::create_directories("/tmp/judge/sandbox");
    }

    void TearDown() override {
        // 清理可能存在的临时文件
        std::filesystem::remove("/tmp/judge/sandbox/test.cpp");
        std::filesystem::remove("/tmp/judge/sandbox/test");
    }
};

TEST_F(SandboxTest, BasicExecution) {
    // 创建沙箱
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 设置命令
    sandbox.setCommand("/bin/echo", {"Hello, World!"});
    
    // 运行命令
    auto result = sandbox.execute();
    
    // 验证结果
    EXPECT_EQ(result.status, ExecutionStatus::Success);
    EXPECT_EQ(result.stdout_output, "Hello, World!\n");
    std::cout << "result.stderr: " << result.stderr_output << std::endl;
    EXPECT_TRUE(result.stderr_output.empty());
}

TEST_F(SandboxTest, ResourceLimits) {
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 设置非常小的内存限制
    ResourceLimits limits;
    limits.memory_bytes = 1 * 1024 * 1024;  // 1MB
    sandbox.setResourceLimits(limits);
    
    // 尝试分配大量内存的程序
    const char* code = R"(
        #include <vector>
        int main() {
            std::vector<char> v(100 * 1024 * 1024);  // 尝试分配100MB
            return 0;
        }
    )";
    
    // 保存代码到临时文件
    std::ofstream("/tmp/judge/sandbox/test.cpp") << code;
    
    // 编译
    sandbox.setCommand("/usr/bin/g++", {"-o", "test", "test.cpp"});
    auto compile_result = sandbox.execute();
    EXPECT_EQ(compile_result.status, ExecutionStatus::Success);
    
    // 运行并期望失败
    sandbox.setCommand("./test", {});
    auto run_result = sandbox.execute();
    
    EXPECT_NE(run_result.status, ExecutionStatus::Success);
    EXPECT_GT(run_result.memory_usage, 0);
}

TEST_F(SandboxTest, CPUTimeLimit) {
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 编译阶段：放宽 CPU 时间
    ResourceLimits compile_limits;
    compile_limits.cpu_time = 5s;
    sandbox.setResourceLimits(compile_limits);
    
    // 创建一个死循环程序
    const char* code = R"(
        #include <iostream>
        #include <cstdio>
        int main() {
            volatile int x = 0;
            long long y = 0;
            while(x < 1000000000) {
                y = (y*2 + x++)/2;
            }
            return 0;
        }
    )";
    
    // 保存代码到临时文件
    std::ofstream("/tmp/judge/sandbox/test.cpp") << code;

    sandbox.setCommand("/usr/bin/g++", {"-o", "test", "test.cpp"});
    auto compile_result = sandbox.execute();
    EXPECT_EQ(compile_result.status, ExecutionStatus::Success);
    // 运行阶段：设置 CPU 时间限制
    ResourceLimits run_limits;
    run_limits.cpu_time = 100ms;
    sandbox.setResourceLimits(run_limits);  // 设置运行阶段限制
    
    // 运行并期望超时
    sandbox.setCommand("./test", {});
    auto run_result = sandbox.execute();
    std::cout << "run_result: " << run_result.stdout_output << " | " << run_result.stderr_output << std::endl;
    EXPECT_EQ(run_result.status, ExecutionStatus::TimeLimitExceeded);
    EXPECT_GE(run_result.cpu_usage.count(), 50);  // 放宽限制，只要超过50ms就算成功
}

TEST_F(SandboxTest, FileSystemIsolation) {
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 尝试访问系统关键文件
    sandbox.setCommand("/bin/cat", {"/etc/shadow"});
    auto result = sandbox.execute();
    
    // 在容器中，文件可能不存在或权限被拒绝
    EXPECT_NE(result.status, ExecutionStatus::Success);
}

TEST_F(SandboxTest, NetworkIsolation) {
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 尝试进行网络连接
    sandbox.setCommand("/usr/bin/curl", {"https://www.google.com"});
    auto result = sandbox.execute();
    
    // 在容器中，命令可能不存在或网络被禁用
    EXPECT_NE(result.status, ExecutionStatus::Success);
}

TEST_F(SandboxTest, ProcessIsolation) {
    Sandbox sandbox;
    
    // 设置工作目录
    sandbox.setWorkingDirectory("/tmp/judge/sandbox");
    
    // 尝试执行系统命令
    const char* code = R"(
        #include <iostream>
        int main() {
            system("ps aux");
            return 0;
        }
    )";
    
    // 保存代码到临时文件
    std::ofstream("/tmp/judge/sandbox/test.cpp") << code;
    
    // 编译
    sandbox.setCommand("/usr/bin/g++", {"-o", "test", "test.cpp"});
    auto compile_result = sandbox.execute();
    EXPECT_EQ(compile_result.status, ExecutionStatus::Success);
    
    // 运行并期望失败
    sandbox.setCommand("./test", {});
    auto run_result = sandbox.execute();
    std::cout << "run_result: " << run_result.stdout_output << " | " << run_result.stderr_output << std::endl;
    // 在容器中，命令可能不存在或被禁用
    EXPECT_NE(run_result.status, ExecutionStatus::Success);
} 